#include "trm/mincut_tracker.h"
#include "mincut_segmenter.h"
#include <imglib/bmpmorph.h>
#include <imglib/bmp.h>
#include <imglib/tga.h>
#include <imglib/time_median.h>
#include <xtd/string>
#include <mswin/seq_name_factory.h>

namespace trm {

namespace mincut {

using std::clog;
using std::endl;
using std::flush;

void tracker::init(size_extent const & dim, uint duration, params const & p, bool overbuild_obligatory_masks)
{
	params_ = p;
	overbuild_obligatory_masks_ = overbuild_obligatory_masks;
	lowlevel_max_motion_ = params_.max_motion * (1 << params_.pyramid_levels);
	dim_ = dim;
	masks_.resize(duration);
	gray_masks_.resize(duration);
	for (uint i = 0; i < duration; ++i) {
		masks_[i].foreground.resize(dim);
		masks_[i].foreground.fill(false);
		masks_[i].background.resize(dim);
		masks_[i].background.fill(false);
	}
}

void tracker::set_obligatory_foreground_pixels(uint frame, mask m)
{
	assert (frame < duration());

	int c = m.set_bit_count();
	if (c == m.size())
		throw std::runtime_error((xtd::string("provided foreground mask for the frame #") += frame) += "has all pixels turned on, which is prohibited");
	if (c == 0)
		return;

	img::combine(masks_[frame].foreground, m, img::bit_or());
	if (!overbuild_obligatory_masks_)
		return;

	for (uint d = 1; d <= frame || frame + d < duration(); ++d) {
		img::erode8(m, lowlevel_max_motion_);
		if (m.set_bit_count() == 0)
			break;
		if (d <= frame)
			img::combine(masks_[frame - d].foreground, m, img::bit_or());
		if (frame + d < duration())
			img::combine(masks_[frame + d].foreground, m, img::bit_or());
	}
}

void tracker::set_obligatory_background_pixels(uint frame, mask m)
{
	assert (frame < duration());

	int c = m.set_bit_count();
	if (c == m.size())
		throw std::runtime_error((xtd::string("provided background mask for the frame #") += frame) += "has all pixels turned on, which is prohibited");
	if (c == 0)
		return;

	img::combine(masks_[frame].background, m, img::bit_or());
	if (!overbuild_obligatory_masks_)
		return;

	for (uint d = 1; d <= frame || frame + d < duration(); ++d) {
		img::erode8(m, lowlevel_max_motion_);
		if (m.set_bit_count() == 0)
			break;
		if (d <= frame)
			img::combine(masks_[frame - d].background, m, img::bit_or());
		if (frame + d < duration())
			img::combine(masks_[frame + d].background, m, img::bit_or());
	}
}

void tracker::set_exact_mask(uint frame, mask const & m)
{
	set_obligatory_foreground_pixels(frame, m);
	mask inv(m);
	inv.inverse();
	set_obligatory_background_pixels(frame, inv);
}

//defined in mincut.cpp
void segment_core(
	byte_image const & image,
	mask const & obligatory_foreground_pixels,
	mask const & obligatory_background_pixels,
	mask * foreground_only,
	mask * background_only,
	params const & p,
	gray_mask * gmask);

/// returns a range [first, last) with at most max_elem that lies inside [0, size) 
/// and is symmetrical relative to given frame
void sym_bounds(uint frame, uint max_elem, uint size, uint & first, uint & last)
{
	assert (max_elem >= 1);
	assert (frame < size);
	if (size < max_elem)
		max_elem = size; ///< do not allow going out both sized of the range
	uint l = max_elem / 2;
	uint r = max_elem - l;
	if (frame < l) {
		first = 0;
		last = 2 * frame + 1;
	}
	else if (frame + r > size) {
		first = 2 * frame - size + 1;
		last = size;
	}
	else {
		first = frame - l;
		last  = frame + r;
	}
	assert (last - first <= max_elem);
}

template <typename T>
class pyramid_bottom
{
public:
	pyramid_bottom(const std::vector< img::pyramid<T> > & pp) : p(pp) { }
	size_t size() const
		{ return p.size(); }
	T const & operator [] (uint n) const
		{ return p[n][0]; }
private:
	const std::vector< img::pyramid<T> > & p;
};

void upsample(
	img::pyramid<mask> & fmask_pyramid,
	img::pyramid<mask> & bmask_pyramid,
	uint level, ///< upsample from level to level-1
	uint max_motion)
{
	size_extent updim = fmask_pyramid[level-1].dim();
	mask f(updim), b(updim);

	f.fill(false);
	img::upsample(fmask_pyramid[level], f, 2);
	img::erode4(f, max_motion);
	img::combine(f, bmask_pyramid[level-1], img::bit_and_not());

	b.fill(false);
	img::upsample(bmask_pyramid[level], b, 2);
	img::erode4(b, max_motion);
	img::combine(b, fmask_pyramid[level-1], img::bit_and_not());

	img::combine(fmask_pyramid[level-1], f, img::bit_or());
	img::combine(bmask_pyramid[level-1], b, img::bit_or());
}

void tracker::solve(icache::footage_sptr const & footage, uint start_frame,
	uint median_len,
	std::string const & out_foremasks_pattern, 
	std::string const & out_backmasks_pattern,
	std::string const & out_graymasks_pattern,
	std::string const & out_alphaimgs_pattern)
{
	//assert (params_.interframe_pyramid_levels > 0); //the user must supply both initial masks for every frame

	//check masks first of all
	size_t ec0 = 0, ec1 = 0;
	for (uint i = 0; i < duration(); ++i) {
		mask m(masks_[i].foreground);
		img::combine(m, masks_[i].background, img::bit_and());
		if (m.set_bit_count() > 0) {
			//img::bmp::write1bit("bad_foremask.bmp", masks_[i].foreground);
			//img::bmp::write1bit("bad_backmask.bmp", masks_[i].background);
			//throw std::runtime_error(
			//	xtd::string(xtd::string("Obligatory foreground and background masks for the frame #") += i) += 
			//		" happen to have common turned on pixels; the masks are saved as bad_foremask.bmp and bad_backmask.bmp.");

			//set this area as ambiguous for both masks
			img::combine(masks_[i].foreground, m, img::bit_and_not());
			img::combine(masks_[i].background, m, img::bit_and_not());
		}
		ec0 += masks_[i].background.set_bit_count();
		ec1 += masks_[i].foreground.set_bit_count();
	}

	if (!params_.silent) {
		clog << 
			"the count of obligatory foreground pixels in all the frames: " << ec1 << "\n"
			"the count of obligatory background pixels in all the frames: " << ec0 << endl;
	}
	if (ec0 == 0)
		throw std::runtime_error("obligatory background pixels mask can't be empty for all the frames");
	if (ec1 == 0)
		throw std::runtime_error("obligatory foreground pixels mask can't be empty for all the frames");

	std::vector< img::pyramid<mask> > fmask_pyramid(duration());
	std::vector< img::pyramid<mask> > bmask_pyramid(duration());

	for (uint i = 0; i < duration(); ++i) {
		fmask_pyramid[i].reset(masks_[i].foreground, params_.pyramid_levels, img::not_copy_original);
		bmask_pyramid[i].reset(masks_[i].background, params_.pyramid_levels, img::not_copy_original);
	}

	bool need_graymasks = !out_graymasks_pattern.empty() || !out_alphaimgs_pattern.empty();
	bool fully_processed = false;

  uint level;
	for (level = params_.pyramid_levels; level + params_.interframe_pyramid_levels > params_.pyramid_levels; --level) 
	{
		if (!params_.silent)
			clog << "\nLevel " << level << endl;

		solve_core(footage, start_frame, level, fmask_pyramid, bmask_pyramid, need_graymasks);

		if (level == 0) {
			fully_processed = true;
			break;
		}

		for (uint i = 0; i < duration(); ++i) 
			upsample(fmask_pyramid[i], bmask_pyramid[i], level, params_.max_motion);
	}

	if (!fully_processed) 
	{
		//process frames on lower levels of pyramid alone (it's faster and sometimes better)
		for (uint i = 0; i < duration(); ++i) 
		{
			icache::cached_image_sptr pc = footage->get(i + start_frame, icache::forward);
			img::pyramid<byte_image> image_pyramid(pc->im, level, img::not_copy_original);

			for (uint l = level; ; --l) 
			{
				segment_core(image_pyramid[l], 
					fmask_pyramid[i][l], bmask_pyramid[i][l], 
					&fmask_pyramid[i][l], &bmask_pyramid[i][l], 
					params_, (need_graymasks && l == 0) ? &gray_masks_[i] : 0);

				if (l == 0)
					break;

				upsample(fmask_pyramid[i], bmask_pyramid[i], l, params_.max_motion);
			}
		}
	}

	{
		mswin::timing_mark mark(params_.timing_tree, "writing out masks");
		if (!out_foremasks_pattern.empty()) {
			mswin::seq_name_factory fac(out_foremasks_pattern);
			fac.set_current(footage->first_frame_num() + start_frame);
			for (uint i = 0; i < duration(); ++i) {
				uint first, last;
				sym_bounds(i, median_len, duration(), first, last);
				mask m(dim_);
				img::bit_time_median(pyramid_bottom<mask>(fmask_pyramid), first, last, m);
				img::bmp::write1bit(fac.get_next().c_str(), m);
			}
		}
		if (!out_backmasks_pattern.empty()) {
			mswin::seq_name_factory fac(out_backmasks_pattern);
			fac.set_current(footage->first_frame_num() + start_frame);
			for (uint i = 0; i < duration(); ++i) {
				uint first, last;
				sym_bounds(i, median_len, duration(), first, last);
				mask m(dim_);
				img::bit_time_median(pyramid_bottom<mask>(bmask_pyramid), first, last, m);
				img::bmp::write1bit(fac.get_next().c_str(), m);
			}
		}
		if (need_graymasks && !gray_masks_[0].empty()) {
			mswin::seq_name_factory fac(out_graymasks_pattern);
			mswin::seq_name_factory fac1(out_alphaimgs_pattern);
			fac.set_current(footage->first_frame_num() + start_frame);
			fac1.set_current(footage->first_frame_num() + start_frame);
			for (uint i = 0; i < duration(); ++i) {
				uint first, last;
				sym_bounds(i, median_len, duration(), first, last);
				gray_mask m(dim_);
				img::time_median(gray_masks_, first, last, m);
				if (!out_graymasks_pattern.empty())
					img::bmp::write8bit(fac.get_next().c_str(), m);
				if (!out_alphaimgs_pattern.empty()) {
					byte_image im = footage->get(i, icache::forward)->im;
					img::projection(im, &byte_color::a) = m;
					img::tga::write32bit(fac1.get_next().c_str(), im);
				}
			}
		}
	}
}

void antialise(
	segmenter & s, 
	std::vector<segmenter::vmap_type> const & vmaps,
	params const & p,
	std::vector<gray_mask> & gmasks)
{
	assert (p.gray_levels_exp <= 6);
	uint dur = vmaps.size();
	for (uint f = 0; f < dur; ++f) {
		gmasks[f].resize(vmaps[f].dim());
		img::fill(gmasks[f], 0);
	}
	for (uint l = 0;; ++l) {
		for (uint f = 0; f < dur; ++f)
			s.inc_frame_gray_mask(vmaps[f], gmasks[f]);
		if (l == p.gray_levels_exp)
			break;
		if (!p.silent)
			clog << "computing gray level " << (l+1) << "..." << endl;
		s.double_capacity_of_saturated_edges();
		s.solve();
	}

	uint m = (1 << (p.gray_levels_exp+2)) - 1;
	//table for converting values in gmask from [0, m) to [0, 256)
	std::vector<unsigned char> table(m);
	for (uint i = 0; i < m; ++i)
		table[i] = (255 * i) / (m - 1);
	for (uint f = 0; f < dur; ++f) {
		for (gray_mask::iterator git = gmasks[f].begin(); git != gmasks[f].end(); ++git) {
			assert (*git < m);
			*git = table[*git];
		}
	}
}

void tracker::solve_core(icache::footage_sptr const & footage, uint start_frame,
	uint level,
	std::vector< img::pyramid<mask> > & fmask_pyramid,
	std::vector< img::pyramid<mask> > & bmask_pyramid,
	bool need_greymasks) ///< only for the level 0
{
	assert (fmask_pyramid.size() == duration());
	assert (bmask_pyramid.size() == duration());

	segmenter s(params_.timing_tree, params_.silent);

	std::vector<segmenter::vmap_type> vmaps(duration());

	if (!params_.silent)
		clog << "Building graph.\n"
			"creating vertices in each of frames" << flush;

	for (uint i = 0; i < duration(); ++i) {
		s.create_vmap(fmask_pyramid[i][level], bmask_pyramid[i][level], vmaps[i]);
		if (!params_.silent)
			clog << '.' << flush;
	}

	if (!params_.silent)
		clog << "\ntotal vertices in graph: " << s.num_vertices() << "\n"
			"creating edges between vertices" << endl;

	// create edges between pixels according to their similarity

	uint scale = 1 << level;
	byte_image prev, curr, next;
	load_image(footage, start_frame, scale, curr);
	for (uint i = 0; i < duration(); ++i) {
		assert (!curr.empty());
		//edges inside the current frame
		s.establish_links(
			curr, curr,
			vmaps[i], vmaps[i],
			params_.hw, params_.nn, params_.aspect);
		if (i > 0) {
			assert (!prev.empty());
			//edges "from" the current frame "to" previous one
			s.establish_links(
				curr, prev, 
				vmaps[i], vmaps[i-1], 
				params_.interframe_hw, params_.interframe_nn / 2, params_.aspect
			);
		}
		if (i + 1 < duration()) { 
			load_image(footage, 1 + i + start_frame, scale, next);
			//edges "from" the current frame "to" next one
			s.establish_links(
				curr, next, 
				vmaps[i], vmaps[i+1], 
				params_.interframe_hw, (params_.interframe_nn + 1) / 2, params_.aspect
			);
		}
		//if (!params_.silent)
		//	clog << '.' << flush;

		prev.swap(curr);
		curr.swap(next);
		next.clear();
	}

	if (!params_.silent)
		clog << "total directional edges in graph: " << s.num_edges() << 
			"\nDone building.\n" << endl;

	s.solve();

	for (uint i = 0; i < duration(); ++i) {
		s.get_frame_mask(vmaps[i], white_color, fmask_pyramid[i][level]);
		s.get_frame_mask(vmaps[i], black_color, bmask_pyramid[i][level]);
	}
	if (need_greymasks && level == 0)
		antialise(s, vmaps, params_, gray_masks_);
}

void tracker::load_image(icache::footage_sptr const & footage, uint frame, uint scale,
	byte_image & im) const
{
	if (scale > 1)
		img::downsample(footage->get(frame, icache::forward)->im, im, scale);
	else
		im = footage->get(frame, icache::forward)->im;
}

} //mincut

} //param
